"""
运行多个daemon.py的办法

比如用0-3 4个显卡卡运行4个daemon.py
CUDA_VISIBLE_DEVICES=0 python daemon.py
CUDA_VISIBLE_DEVICES=1 python daemon.py
CUDA_VISIBLE_DEVICES=2 python daemon.py
CUDA_VISIBLE_DEVICES=3 python daemon.py

"""

import time
import redis
from util_mongo import get_history
import pymongo as pm
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX

# 连接本机6379端口上的redis的第0个数据库
rdb = redis.Redis('127.0.0.1', 6379, 0)

# 连接Mongodb
mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
mdb = mongo[MONGODB_NAME]

# 定义映射
fake_chat_dict = {
    0: '十全十美！',
    1: '一把钢枪交给我！',
    2: '二话不说为祖国！',
    3: '三山五岳任我走。',
    4: '四海为家。',
    5: '五福同寿。',
    6: '六六大顺！',
    7: '97香港回归。',
    8: '零八奥运。',
    9: '九九归一。',
}


def model_infer(xinput, username):
    """
    NLP LLM模型

    :param xinput: 输入文本
    :return: 根据输入文本的长度，输出相应的文本
    """
    # 获取聊天历史
    xlog = get_history(mdb, username, more_info=False)
    
    # 模型推理
    print('-------------history-----------------')
    for i, (xin, xout) in enumerate(xlog):
        print(i, '>>>>', xin)
        print(i, '<<<<', xout)
    print('-------------this turn---------------')
    print('>>>>', '>>>>', xinput)

    # 根据输入文本的长度，输出相应的文本
    time.sleep(1.0)
    xlen = len(xinput)
    xoutput = f'您说了{xlen}个字符。（{xinput[:5]}……）' + fake_chat_dict[xlen % 10] + f'(多轮对话之前轮数{len(xlog)})'

    print('<<<<', '<<<<', xoutput)
    
    return xoutput


if '__main__' == __name__:

    print('------------------READY---------------------------')

    # 后台一直运行
    while True:

        # 从管道拿uuid
        xuuid = rdb.rpop('queue')

        if xuuid is None:
            # 没有数据则等1ms后重试
            time.sleep(0.001)
            continue
        xuuid = xuuid.decode('utf8')
        print('UUID:', xuuid)

        # 拿输入
        xinput = rdb.hget('uuid2input', xuuid)
        if xinput is None:
            # 为了健壮性
            time.sleep(0.001)
            continue
        xinput = xinput.decode('utf8')
        print('input:', xinput)
        
        # 拿用户名
        username = rdb.hget('uuid2username', xuuid)
        if username is None:
            # 为了健壮性
            time.sleep(0.001)
            continue
        username = username.decode('utf8')
        print('username:', username)

        # 模型推理
        xoutput = model_infer(xinput, username)
        print('output:', xoutput)

        # 把输出放回
        rdb.hset('uuid2output', xuuid, xoutput.encode('utf8'))
